package org.skipper.dubbo.spring.starter.helper;

import lombok.extern.log4j.Log4j2;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.apache.logging.log4j.ThreadContext;
import org.skipper.dubbo.spring.starter.model.RpcTraceContext;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

import static org.apache.dubbo.common.constants.CommonConstants.CONSUMER;
import static org.apache.dubbo.common.constants.CommonConstants.PROVIDER;

/**
 * Created with IntelliJ IDEA.
 * Description:
 *
 * @author shoujian
 * Date: 2022-02-15
 * Time: 2:15 PM
 */
@Activate(group = {PROVIDER})
@Log4j2
public class DubboProviderContextFilter implements Filter {

    private static final String TRACE_ID = "traceId";

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        Map<String, Object> objectAttachments = RpcContext.getContext().getObjectAttachments();

        String traceId = UUID.randomUUID().toString();

        if (objectAttachments.containsKey("traceId")) {
            traceId = objectAttachments.get("traceId").toString();
        } else {
            objectAttachments.put("traceId", traceId);
        }
        // 存入 Log4j2 ThreadContext
        ThreadContext.put(TRACE_ID, traceId);

        RpcTraceContext.setAttachments(objectAttachments);

        try {
            return invoker.invoke(invocation);
        } finally {
            RpcTraceContext.clear();
        }
    }
}
